#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
PostgreSQL用户管理模块
提供用户创建、删除和密码管理功能
"""

import os
from typing import Optional

from ...core.utils import (
    logger, 
    run_command, 
    get_input_with_default, 
    confirm_action,
    is_windows
)
from ...core.ui_utils import console
from ...core import config

class PostgresqlUser:
    """PostgreSQL用户管理类"""
    
    def create_database_and_user(self):
        """创建新的数据库和用户"""
        db_name = get_input_with_default("输入要创建的数据库名称", "")
        db_user = get_input_with_default("输入要创建的数据库用户名", "")
        db_password = get_input_with_default("输入创建用户的密码", "")
        
        if not db_name or not db_user or not db_password:
            logger.error("数据库名、用户名和密码不能为空")
            return
        
        # 检查数据库是否已经存在
        try:
            check_db_command = f"SELECT 1 FROM pg_database WHERE datname = '{db_name}';"
            db_exists = run_command(
                ["sudo", "-u", "postgres", "psql", "-tAc", check_db_command],
                capture_output=True
            ).stdout.strip()
            
            if db_exists == '1':
                logger.error(f"数据库 '{db_name}' 已存在")
                return
        except Exception as e:
            logger.error(f"检查数据库时出错: {e}")
            return
        
        # 检查用户是否已经存在
        try:
            check_user_command = f"SELECT 1 FROM pg_roles WHERE rolname = '{db_user}';"
            user_exists = run_command(
                ["sudo", "-u", "postgres", "psql", "-tAc", check_user_command],
                capture_output=True
            ).stdout.strip()
            
            if user_exists == '1':
                logger.error(f"用户 '{db_user}' 已存在")
                return
        except Exception as e:
            logger.error(f"检查用户时出错: {e}")
            return
        
        # 创建用户和数据库
        try:
            logger.info(f"正在创建用户 {db_user}")
            run_command(["sudo", "-u", "postgres", "createuser", db_user])
            
            logger.info(f"正在设置用户 {db_user} 密码")
            run_command([
                "sudo", "-u", "postgres", "psql", 
                "-c", f"ALTER USER {db_user} WITH PASSWORD '{db_password}';"
            ])
            
            logger.info(f"正在创建数据库 {db_name}")
            run_command(["sudo", "-u", "postgres", "createdb", db_name])
            
            logger.info(f"正在授权用户 {db_user} 访问数据库 {db_name}")
            run_command([
                "sudo", "-u", "postgres", "psql", 
                "-c", f"GRANT ALL PRIVILEGES ON DATABASE {db_name} TO {db_user};"
            ])
            
            logger.info(f"数据库 '{db_name}' 和用户 '{db_user}' 已成功创建")
        except Exception as e:
            logger.error(f"创建数据库和用户时出错: {e}")
    
    def change_postgres_password(self):
        """修改PostgreSQL用户密码"""
        username = get_input_with_default("输入PostgreSQL用户名", "postgres")
        new_password = get_input_with_default(f"输入用户 '{username}' 的新密码", "")
        
        if not new_password:
            logger.error("密码不能为空")
            return
        
        try:
            change_password_command = f"ALTER USER {username} WITH PASSWORD '{new_password}';"
            run_command(["sudo", "-u", "postgres", "psql", "-c", change_password_command])
            logger.info(f"用户 '{username}' 的密码已成功修改")
        except Exception as e:
            logger.error(f"修改用户 '{username}' 密码时出错: {e}")
    
    def show_all_users(self):
        """显示所有数据库用户"""
        logger.info("数据库用户列表:")
        
        try:
            run_command(["sudo", "-u", "postgres", "psql", "-c", "SELECT rolname, rolsuper FROM pg_roles;"])
        except Exception as e:
            logger.error(f"获取用户列表时出错: {e}")
    
    def delete_user(self):
        """删除指定的数据库用户"""
        username = get_input_with_default("输入要删除的数据库用户名", "")
        
        if not username:
            logger.error("用户名不能为空")
            return
        
        if username == "postgres":
            logger.error("不能删除系统用户 'postgres'")
            return
        
        # 检查用户是否存在
        try:
            check_user_command = f"SELECT 1 FROM pg_roles WHERE rolname = '{username}';"
            user_exists = run_command(
                ["sudo", "-u", "postgres", "psql", "-tAc", check_user_command],
                capture_output=True
            ).stdout.strip()
            
            if user_exists != '1':
                logger.error(f"用户 '{username}' 不存在")
                return
        except Exception as e:
            logger.error(f"检查用户时出错: {e}")
            return
        
        # 删除用户
        if confirm_action(f"确定要删除用户 '{username}' 吗？", False):
            try:
                run_command(["sudo", "-u", "postgres", "psql", "-c", f"DROP ROLE {username};"])
                logger.info(f"用户 '{username}' 已成功删除")
            except Exception as e:
                logger.error(f"删除用户 '{username}' 时出错: {e}")
    
    def create_replicator_user(self):
        """创建用于流复制的专用用户"""
        username = config.DEFAULT_REPLICATION_USER
        password = get_input_with_default(f"输入流复制用户 '{username}' 的密码", "")
        
        if not password:
            logger.error("密码不能为空")
            return
        
        # 检查用户是否已存在
        try:
            check_user_command = f"SELECT 1 FROM pg_roles WHERE rolname = '{username}';"
            user_exists = run_command(
                ["sudo", "-u", "postgres", "psql", "-tAc", check_user_command],
                capture_output=True
            ).stdout.strip()
            
            if user_exists == '1':
                logger.error(f"用户 '{username}' 已存在")
                return
        except Exception as e:
            logger.error(f"检查用户时出错: {e}")
            return
        
        # 创建复制用户
        try:
            create_user_command = f"CREATE ROLE {username} WITH REPLICATION LOGIN PASSWORD '{password}';"
            run_command(["sudo", "-u", "postgres", "psql", "-c", create_user_command])
            logger.info(f"流复制用户 '{username}' 已成功创建")
        except Exception as e:
            logger.error(f"创建流复制用户时出错: {e}")
            
    def create_pgpass_file(self):
        """为当前客户端创建免密登录数据库配置文件"""
        hostname = get_input_with_default("请输入数据库链接IP", "localhost")
        port = get_input_with_default("请输入数据库端口", "5432")
        database = get_input_with_default("请输入要登录的数据库名", "postgres")
        username = get_input_with_default("请输入数据库用户名", "postgres")
        password = get_input_with_default("请输入数据库密码", "")
        
        if not password:
            logger.error("密码不能为空")
            return
            
        # 创建.pgpass文件的内容
        pgpass_entry = f"{hostname}:{port}:{database}:{username}:{password}\n"
        
        # 获取当前用户和postgres用户的主目录路径
        current_user_home = os.path.expanduser("~")
        
        if is_windows():
            # Windows上用户目录
            pgpass_paths = [os.path.join(current_user_home, "AppData\\Roaming\\postgresql\\pgpass.conf")]
        else:
            # Linux用户目录和postgres用户目录
            root_home_dir = "/root"
            postgres_home_dir = "/var/lib/postgresql"
            pgpass_paths = [
                os.path.join(current_user_home, ".pgpass"),
                os.path.join(root_home_dir, ".pgpass"),
                os.path.join(postgres_home_dir, ".pgpass")
            ]
            
        try:
            for pgpass_path in pgpass_paths:
                # 确保目录存在
                pgpass_dir = os.path.dirname(pgpass_path)
                if not os.path.exists(pgpass_dir):
                    logger.info(f"目录 {pgpass_dir} 不存在，跳过")
                    continue
                    
                if os.path.exists(pgpass_path):
                    # 如果文件存在，检查是否已经包含条目
                    with open(pgpass_path, 'r') as pgpass_file:
                        content = pgpass_file.read()
                    if pgpass_entry not in content:
                        # 如果条目不存在，追加新的条目
                        with open(pgpass_path, 'a') as pgpass_file:
                            pgpass_file.write(pgpass_entry)
                        logger.info(f"条目已添加到 {pgpass_path}")
                    else:
                        logger.info(f"条目已存在于 {pgpass_path}")
                else:
                    # 如果文件不存在，创建文件并写入条目
                    with open(pgpass_path, 'w') as pgpass_file:
                        pgpass_file.write(pgpass_entry)
                    logger.info(f".pgpass文件创建于 {pgpass_path}")
                    
                # 设置文件权限为 0600
                os.chmod(pgpass_path, 0o600)
                logger.info(f"文件权限设置为600: {pgpass_path}")
                
            # 确保 postgres 用户对其主目录中的 .pgpass 文件有权限
            if not is_windows():
                postgres_pgpass = "/var/lib/postgresql/.pgpass"
                if os.path.exists(postgres_pgpass):
                    run_command(['sudo', 'chown', 'postgres:postgres', postgres_pgpass])
                    
            logger.info("配置文件创建完成")
        except Exception as e:
            logger.error(f"创建 .pgpass 文件时出错: {e}") 

    def drop_database(self):
        """删除指定数据库"""
        db_name = get_input_with_default("输入要删除的数据库名称", "")
        
        if not db_name:
            logger.error("数据库名不能为空")
            return
        
        if db_name.lower() in ["postgres", "template0", "template1"]:
            logger.error(f"不能删除系统数据库 '{db_name}'")
            return
        
        # 检查数据库是否存在
        try:
            check_db_command = f"SELECT 1 FROM pg_database WHERE datname = '{db_name}';"
            db_exists = run_command(
                ["sudo", "-u", "postgres", "psql", "-tAc", check_db_command],
                capture_output=True
            ).stdout.strip()
            
            if db_exists != '1':
                logger.error(f"数据库 '{db_name}' 不存在")
                return
        except Exception as e:
            logger.error(f"检查数据库时出错: {e}")
            return
        
        # 确认并删除数据库
        if confirm_action(f"确定要删除数据库 '{db_name}' 吗？此操作将永久删除所有数据", False):
            try:
                # 断开所有连接
                drop_connections_command = f"""
                SELECT pg_terminate_backend(pg_stat_activity.pid)
                FROM pg_stat_activity
                WHERE pg_stat_activity.datname = '{db_name}'
                AND pid <> pg_backend_pid();
                """
                logger.info(f"正在断开所有到数据库 '{db_name}' 的连接...")
                run_command(["sudo", "-u", "postgres", "psql", "-c", drop_connections_command])
                
                # 删除数据库
                logger.info(f"正在删除数据库 '{db_name}'...")
                run_command(["sudo", "-u", "postgres", "dropdb", db_name])
                logger.info(f"数据库 '{db_name}' 已成功删除")
            except Exception as e:
                logger.error(f"删除数据库 '{db_name}' 时出错: {e}")
    
    def truncate_database(self):
        """清空指定数据库的所有表和数据"""
        db_name = get_input_with_default("输入要清空的数据库名称", "")
        
        if not db_name:
            logger.error("数据库名不能为空")
            return
        
        if db_name.lower() in ["postgres", "template0", "template1"]:
            logger.error(f"不能清空系统数据库 '{db_name}'")
            return
        
        # 检查数据库是否存在
        try:
            check_db_command = f"SELECT 1 FROM pg_database WHERE datname = '{db_name}';"
            db_exists = run_command(
                ["sudo", "-u", "postgres", "psql", "-tAc", check_db_command],
                capture_output=True
            ).stdout.strip()
            
            if db_exists != '1':
                logger.error(f"数据库 '{db_name}' 不存在")
                return
        except Exception as e:
            logger.error(f"检查数据库时出错: {e}")
            return
        
        # 确认并清空数据库
        if confirm_action(f"确定要清空数据库 '{db_name}' 的所有表和数据吗？此操作将永久删除所有数据", False):
            try:
                # 先断开所有连接
                drop_connections_command = f"""
                SELECT pg_terminate_backend(pg_stat_activity.pid)
                FROM pg_stat_activity
                WHERE pg_stat_activity.datname = '{db_name}'
                AND pid <> pg_backend_pid();
                """
                logger.info(f"正在断开所有到数据库 '{db_name}' 的连接...")
                run_command(["sudo", "-u", "postgres", "psql", "-c", drop_connections_command])
                
                # 获取所有表名，并在单个事务中截断所有表
                get_tables_command = f"""
                SELECT 'TRUNCATE TABLE ' || string_agg(quote_ident(tablename), ', ') || ' CASCADE;'
                FROM pg_tables
                WHERE schemaname = 'public';
                """
                
                logger.info(f"正在获取数据库 '{db_name}' 的所有表...")
                truncate_command = run_command(
                    ["sudo", "-u", "postgres", "psql", "-d", db_name, "-tAc", get_tables_command],
                    capture_output=True
                ).stdout.strip()
                
                if truncate_command and truncate_command != "TRUNCATE TABLE  CASCADE;":
                    logger.info(f"正在清空数据库 '{db_name}' 的所有表...")
                    run_command(["sudo", "-u", "postgres", "psql", "-d", db_name, "-c", truncate_command])
                    logger.info(f"数据库 '{db_name}' 的所有表已成功清空")
                else:
                    logger.info(f"数据库 '{db_name}' 中没有可清空的表")
                    
                # 删除序列
                reset_sequences_command = f"""
                DO $$
                DECLARE
                    seq record;
                BEGIN
                    FOR seq IN SELECT relname FROM pg_class WHERE relkind = 'S' AND relnamespace = 'public'::regnamespace LOOP
                        EXECUTE 'ALTER SEQUENCE ' || quote_ident(seq.relname) || ' RESTART WITH 1;';
                    END LOOP;
                END $$;
                """
                
                logger.info(f"正在重置数据库 '{db_name}' 的所有序列...")
                run_command(["sudo", "-u", "postgres", "psql", "-d", db_name, "-c", reset_sequences_command])
                logger.info(f"数据库 '{db_name}' 的所有序列已重置")
                
            except Exception as e:
                logger.error(f"清空数据库 '{db_name}' 时出错: {e}")

    def drop_all_tables(self):
        """删除指定数据库中的所有表"""
        db_name = get_input_with_default("输入要清空所有表的数据库名称", "")
        
        if not db_name:
            logger.error("数据库名不能为空")
            return
        
        if db_name.lower() in ["postgres", "template0", "template1"]:
            logger.error(f"不能清空系统数据库 '{db_name}'")
            return
        
        # 检查数据库是否存在
        try:
            check_db_command = f"SELECT 1 FROM pg_database WHERE datname = '{db_name}';"
            db_exists = run_command(
                ["sudo", "-u", "postgres", "psql", "-tAc", check_db_command],
                capture_output=True
            ).stdout.strip()
            
            if db_exists != '1':
                logger.error(f"数据库 '{db_name}' 不存在")
                return
        except Exception as e:
            logger.error(f"检查数据库时出错: {e}")
            return
        
        # 确认并删除所有表
        if confirm_action(f"确定要删除数据库 '{db_name}' 中的所有表吗？此操作将永久删除所有表结构", False):
            try:
                # 先断开所有连接
                drop_connections_command = f"""
                SELECT pg_terminate_backend(pg_stat_activity.pid)
                FROM pg_stat_activity
                WHERE pg_stat_activity.datname = '{db_name}'
                AND pid <> pg_backend_pid();
                """
                logger.info(f"正在断开所有到数据库 '{db_name}' 的连接...")
                run_command(["sudo", "-u", "postgres", "psql", "-c", drop_connections_command])
                
                # 获取所有表名
                get_tables_command = f"""
                SELECT string_agg(quote_ident(tablename), ', ')
                FROM pg_tables
                WHERE schemaname = 'public';
                """
                
                logger.info(f"正在获取数据库 '{db_name}' 的所有表...")
                tables = run_command(
                    ["sudo", "-u", "postgres", "psql", "-d", db_name, "-tAc", get_tables_command],
                    capture_output=True
                ).stdout.strip()
                
                if tables:
                    # 构建删除表的SQL命令
                    drop_tables_command = f"DROP TABLE IF EXISTS {tables} CASCADE;"
                    
                    logger.info(f"正在删除数据库 '{db_name}' 的所有表...")
                    run_command(["sudo", "-u", "postgres", "psql", "-d", db_name, "-c", drop_tables_command])
                    logger.info(f"数据库 '{db_name}' 的所有表已成功删除")
                else:
                    logger.info(f"数据库 '{db_name}' 中没有可删除的表")
                
            except Exception as e:
                logger.error(f"删除数据库 '{db_name}' 的所有表时出错: {e}")

    def truncate_tables(self):
        """清空指定数据库中的指定表的数据"""
        db_name = get_input_with_default("输入数据库名称", "")
        
        if not db_name:
            logger.error("数据库名不能为空")
            return
        
        if db_name.lower() in ["postgres", "template0", "template1"]:
            logger.error(f"不能清空系统数据库 '{db_name}'")
            return
        
        # 检查数据库是否存在
        try:
            check_db_command = f"SELECT 1 FROM pg_database WHERE datname = '{db_name}';"
            db_exists = run_command(
                ["sudo", "-u", "postgres", "psql", "-tAc", check_db_command],
                capture_output=True
            ).stdout.strip()
            
            if db_exists != '1':
                logger.error(f"数据库 '{db_name}' 不存在")
                return
        except Exception as e:
            logger.error(f"检查数据库时出错: {e}")
            return
        
        # 获取所有表名
        try:
            # 先断开所有连接
            drop_connections_command = f"""
            SELECT pg_terminate_backend(pg_stat_activity.pid)
            FROM pg_stat_activity
            WHERE pg_stat_activity.datname = '{db_name}'
            AND pid <> pg_backend_pid();
            """
            logger.info(f"正在断开其他到数据库 '{db_name}' 的连接...")
            run_command(["sudo", "-u", "postgres", "psql", "-c", drop_connections_command])
        
            get_tables_command = f"""
            SELECT tablename
            FROM pg_tables
            WHERE schemaname = 'public'
            ORDER BY tablename;
            """
            
            tables_result = run_command(
                ["sudo", "-u", "postgres", "psql", "-d", db_name, "-tAc", get_tables_command],
                capture_output=True
            ).stdout.strip()
            
            if not tables_result:
                logger.error(f"数据库 '{db_name}' 中没有表")
                return
                
            available_tables = [t.strip() for t in tables_result.split('\n') if t.strip()]
            
            # 显示可用表
            console.print("\n可用的表:", style="bold green")
            for i, table in enumerate(available_tables, 1):
                console.print(f"{i}. {table}")
            
            # 获取用户输入的表名
            table_input = get_input_with_default("请输入要清空的表名（多个表用逗号分隔）", "")
            
            if not table_input:
                logger.error("表名不能为空")
                return
                
            # 处理输入的表名
            selected_tables = [t.strip() for t in table_input.split(',')]
            
            # 验证表是否存在
            invalid_tables = [t for t in selected_tables if t not in available_tables]
            if invalid_tables:
                logger.error(f"以下表不存在: {', '.join(invalid_tables)}")
                return
            
            # 确认并清空表
            if confirm_action(f"确定要清空以下表的数据吗？\n{', '.join(selected_tables)}", False):
                try:
                    # 构建TRUNCATE命令，为表名添加双引号
                    quoted_tables = [f'"{table}"' for table in selected_tables]
                    truncate_command = f"TRUNCATE TABLE {', '.join(quoted_tables)} CASCADE;"
                    
                    logger.info(f"正在清空表...")
                    run_command(["sudo", "-u", "postgres", "psql", "-d", db_name, "-c", truncate_command])
                    
                    # 重置每个表的序列
                    for table in selected_tables:
                        # 获取表的主键列名
                        get_pk_command = f"""
                        SELECT a.attname
                        FROM pg_index i
                        JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
                        WHERE i.indrelid = '"{table}"'::regclass
                        AND i.indisprimary;
                        """
                        
                        pk_column = run_command(
                            ["sudo", "-u", "postgres", "psql", "-d", db_name, "-tAc", get_pk_command],
                            capture_output=True
                        ).stdout.strip()
                        
                        if pk_column:
                            # 获取序列名称
                            get_seq_command = f"""
                            SELECT pg_get_serial_sequence('"{table}"', '{pk_column}');
                            """
                            
                            seq_name = run_command(
                                ["sudo", "-u", "postgres", "psql", "-d", db_name, "-tAc", get_seq_command],
                                capture_output=True
                            ).stdout.strip()
                            
                            if seq_name:
                                # 重置序列
                                reset_seq_command = f"ALTER SEQUENCE {seq_name} RESTART WITH 1;"
                                run_command(["sudo", "-u", "postgres", "psql", "-d", db_name, "-c", reset_seq_command])
                                logger.info(f"表 '{table}' 的序列已重置")
                    
                    logger.info("所有选定的表已成功清空")
                    
                except Exception as e:
                    logger.error(f"清空表时出错: {e}")
                    
        except Exception as e:
            logger.error(f"获取表列表时出错: {e}")

    def sync_sequence(self):
        """同步指定数据库指定表的id序列"""
        db_name = get_input_with_default("输入数据库名称", "")
        table_name = get_input_with_default("输入表名称", "")
        
        if not db_name or not table_name:
            logger.error("数据库名和表名不能为空")
            return
        
        # 检查数据库是否存在
        try:
            check_db_command = f"SELECT 1 FROM pg_database WHERE datname = '{db_name}';"
            db_exists = run_command(
                ["sudo", "-u", "postgres", "psql", "-tAc", check_db_command],
                capture_output=True
            ).stdout.strip()
            
            if db_exists != '1':
                logger.error(f"数据库 '{db_name}' 不存在")
                return
        except Exception as e:
            logger.error(f"检查数据库时出错: {e}")
            return
        
        # 检查表是否存在
        try:
            check_table_command = f"SELECT 1 FROM information_schema.tables WHERE table_schema = 'public' AND table_name = '{table_name}';"
            table_exists = run_command(
                ["sudo", "-u", "postgres", "psql", "-d", db_name, "-tAc", check_table_command],
                capture_output=True
            ).stdout.strip()
            
            if table_exists != '1':
                logger.error(f"表 '{table_name}' 在数据库 '{db_name}' 中不存在")
                return
        except Exception as e:
            logger.error(f"检查表时出错: {e}")
            return
        
        # 获取表的主键列名
        try:
            get_pk_command = f"""
            SELECT a.attname
            FROM pg_index i
            JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
            WHERE i.indrelid = '{table_name}'::regclass
            AND i.indisprimary;
            """
            
            pk_column = run_command(
                ["sudo", "-u", "postgres", "psql", "-d", db_name, "-tAc", get_pk_command],
                capture_output=True
            ).stdout.strip()
            
            if not pk_column:
                logger.error(f"表 '{table_name}' 没有主键")
                return
                
            logger.info(f"表 '{table_name}' 的主键列是 '{pk_column}'")
            
            # 获取序列名称
            get_seq_command = f"""
            SELECT pg_get_serial_sequence('{table_name}', '{pk_column}');
            """
            
            seq_name = run_command(
                ["sudo", "-u", "postgres", "psql", "-d", db_name, "-tAc", get_seq_command],
                capture_output=True
            ).stdout.strip()
            
            if not seq_name:
                logger.error(f"表 '{table_name}' 的主键列 '{pk_column}' 没有关联的序列")
                return
                
            logger.info(f"表 '{table_name}' 的序列是 '{seq_name}'")
            
            # 获取当前最大ID值
            get_max_id_command = f"SELECT COALESCE(MAX({pk_column}), 0) FROM {table_name};"
            
            max_id = run_command(
                ["sudo", "-u", "postgres", "psql", "-d", db_name, "-tAc", get_max_id_command],
                capture_output=True
            ).stdout.strip()
            
            if not max_id.isdigit():
                logger.error(f"获取表 '{table_name}' 的最大ID值时出错")
                return
                
            max_id = int(max_id)
            logger.info(f"表 '{table_name}' 当前最大ID值是 {max_id}")
            
            # 重置序列
            reset_seq_command = f"ALTER SEQUENCE {seq_name} RESTART WITH {max_id + 1};"
            
            logger.info(f"正在重置序列 '{seq_name}' 为 {max_id + 1}...")
            run_command(["sudo", "-u", "postgres", "psql", "-d", db_name, "-c", reset_seq_command])
            logger.info(f"序列 '{seq_name}' 已成功重置")
            
        except Exception as e:
            logger.error(f"同步序列时出错: {e}") 